/**
 * Copyright 2019-2023 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "tensor/image_config_tensor_util.h"
#include "graph/buffer.h"

#include "securec.h"

namespace hiai {
std::shared_ptr<INDTensorBuffer> CreateImageConfigTensorWithSize(uint8_t*& para, size_t paraSize)
{
    NDTensorDesc desc;
    desc.dims = {1, static_cast<int>(paraSize), 1, 1};
    desc.dataType = DataType::UINT8;
    desc.format = Format::NCHW;

    return CreateNDTensorBuffer(desc, para, paraSize);
}

template<typename T>
uint8_t* CreateUnifiedPara(T* paraData, const size_t paraNums, const ImageFormat& imageFormat);

template<typename T>
uint8_t* CreateUnifiedPara(T* paraData, const size_t paraNums, const ImageFormat& imageFormat)
{
    size_t tensorSize = paraNums * sizeof(T) + sizeof(ImageFormat);
    uint8_t* tensorData = reinterpret_cast<uint8_t*>(malloc(tensorSize));

    (void)memcpy_s(tensorData, sizeof(ImageFormat), &imageFormat, sizeof(ImageFormat));

    uint32_t memOffset = sizeof(ImageFormat);
    (void)memcpy_s(tensorData + memOffset, paraNums * sizeof(T), paraData, paraNums * sizeof(T));

    return tensorData;
}

std::shared_ptr<INDTensorBuffer> ImageConfigTensorUtil::CreateImageConfigTensor(MultiCropPara& para)
{
    size_t paraSize = sizeof(CropPara) * para.cropParas.size() + sizeof(para.imageFormat);
    uint8_t* tensorBuffer = CreateUnifiedPara<CropPara>(para.cropParas.data(), para.cropParas.size(), para.imageFormat);

    std::shared_ptr<INDTensorBuffer> paraTensorBuffer = CreateImageConfigTensorWithSize(tensorBuffer, paraSize);
    free(tensorBuffer);

    return paraTensorBuffer;
}

std::shared_ptr<INDTensorBuffer> ImageConfigTensorUtil::CreateImageConfigTensor(MultiResizePara& para)
{
    size_t paraSize = sizeof(ResizePara) * para.resizeParas.size() + sizeof(para.imageFormat);
    uint8_t* tensorBuffer =
        CreateUnifiedPara<ResizePara>(para.resizeParas.data(), para.resizeParas.size(), para.imageFormat);

    std::shared_ptr<INDTensorBuffer> paraTensorBuffer = CreateImageConfigTensorWithSize(tensorBuffer, paraSize);
    free(tensorBuffer);

    return paraTensorBuffer;
}

std::shared_ptr<INDTensorBuffer> ImageConfigTensorUtil::CreateImageConfigTensor(MultiDtcPara& para)
{
    size_t paraSize = sizeof(DtcPara) * para.dtcParas.size() + sizeof(para.imageFormat);
    uint8_t* tensorBuffer = CreateUnifiedPara<DtcPara>(para.dtcParas.data(), para.dtcParas.size(), para.imageFormat);

    std::shared_ptr<INDTensorBuffer> paraTensorBuffer = CreateImageConfigTensorWithSize(tensorBuffer, paraSize);
    free(tensorBuffer);

    return paraTensorBuffer;
}

std::shared_ptr<INDTensorBuffer> ImageConfigTensorUtil::CreateImageConfigTensor(MultiRotatePara& para)
{
    size_t paraSize = sizeof(RotatePara) * para.rotateParas.size() + sizeof(para.imageFormat);
    uint8_t* tensorBuffer =
        CreateUnifiedPara<RotatePara>(para.rotateParas.data(), para.rotateParas.size(), para.imageFormat);

    std::shared_ptr<INDTensorBuffer> paraTensorBuffer = CreateImageConfigTensorWithSize(tensorBuffer, paraSize);
    free(tensorBuffer);

    return paraTensorBuffer;
}

std::shared_ptr<INDTensorBuffer> ImageConfigTensorUtil::CreateImageConfigTensor(MultiPadPara& para)
{
    size_t paraSize = sizeof(PadPara) * para.padParas.size() + sizeof(para.imageFormat);
    uint8_t* tensorBuffer = CreateUnifiedPara<PadPara>(para.padParas.data(), para.padParas.size(), para.imageFormat);

    std::shared_ptr<INDTensorBuffer> paraTensorBuffer = CreateImageConfigTensorWithSize(tensorBuffer, paraSize);
    free(tensorBuffer);

    return paraTensorBuffer;
}

std::shared_ptr<INDTensorBuffer> ImageConfigTensorUtil::CreateImageConfigTensor(ChannelSwapPara& para)
{
    NDTensorDesc desc;
    int channel = static_cast<int>(sizeof(ChannelSwapPara));
    desc.dims = {1, channel, 1, 1};
    desc.dataType = DataType::UINT8;
    desc.format = Format::NCHW;

    return CreateNDTensorBuffer(desc, &para, sizeof(ChannelSwapPara));
}

std::shared_ptr<INDTensorBuffer> ImageConfigTensorUtil::CreateImageConfigTensor(CscPara& para)
{
    NDTensorDesc desc;
    int channel = static_cast<int>(sizeof(CscPara));
    desc.dims = {1, channel, 1, 1};
    desc.dataType = DataType::UINT8;
    desc.format = Format::NCHW;

    return CreateNDTensorBuffer(desc, &para, sizeof(CscPara));
}
} // namespace hiai